/***************************************
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@!!!@@@@@@@@@@@@!!!!@@@@@@@@@@@%!!!@@@@@@@@@@@!!!!!!!!!!@@
@@!!!!!@@@@@@@@@%%!!!!!@@@@@@@@@%%!!!!@@@@@@@@@!!!!!!!!!!@@
@@@!!!!!@@@@@@@%%%%!!!!!@@@@@@@%%%!!!!!@@@@@@@@@@@@@!!!!!@@
@@@@!!!!!@@@@@@%%%@@!!!!!@@@@@%%%%@!!!!!@@@@@@@@@@@@!!!!!@@
@@@@@!!!!!@@@@%%%@@@@!!!!@@@@%%%%@@@!!!!!@@@@@@@@@@@!!!!!@@
@@@@@@!!!!!@@%%%%@@@@!!!!!@@@%%%@@@@@!!!!!@@@@@@@@@@!!!!!@@
@@@@@@!!!!!@%%%%@@@@@@!!!!!@%%%$@@@@@@!!!!@@@@@@@@@@!!!!!@@
@@@@@@@!!!!!%%%@@@@@@@@!!!!!%%%@@@@@@@@!!!!@@@@@@@@!!!!!!@@
@@@@@@@@!!!!!%@@@@@@@@@@!!!!!%@@@@@@@@@!!!!!!!!!!!!!!!!!!@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

WMJ战队深度学习装甲板目标检测类
支持的模型：          yolox改四点模型
时间：               2024.01.31
***************************************/
#pragma once
#include "armor/Armor.hpp"
#include "libbase/common.h"
#include "InferenceBackend.hpp"
namespace wmj
{
    struct bbox_t {
        cv::Point2f pts[4];    // [pt0, pt1, pt2, pt3] 四点组，顺序为：左上↖，左下↙，右下↘，右上↗
        float conf;            // 综合置信度
        cv::Rect_<float> rect; // 外接轮廓
        float color_conf;      // 颜色置信度 
        float type_conf;
        int color_id; // 0: red, 1: blue, 2: gray
        int type_id;
        int tag_id;
    };

    struct GridAndStride
    {
        int grid0;
        int grid1;
        int stride;
    };

    struct DetectorParams
    {
        int enemyColor;
        // int DEVICE;
        std::string DEVICE;
        float lightRatio;
        float nmsThreshold;
        float confThreshold;
        std::string modelName;
        std::string modelFormat;
        std::string modelType;
        std::string inputPrecision;
        int nColor;
        int nType;
        int nTag;
        bool debug;
    };

    class DeepDetector
    {
        /*********************深度学习检测器主类*********************/
        public:
        /**
         * @brief Construct a new Deep Detector object 无参构造，默认读取 DEEP_CFG 文件中参数
         * 
         */
        DeepDetector();
        /**
         * @brief Construct a new Deep Detector object 带参构造，读取参数类中参数
         * 
         * @param params 参数类
         */
        DeepDetector(const DetectorParams &params);
        /**
         * @brief Destroy the Deep Detector object 析构函数
         * 
         */
        ~DeepDetector();

        // DeepDetector(const std::string xml, const std::string bin);                // 带参构造指定模型文件
        /**
         * @brief 类初始化函数，初始化类成员
         * 
         */
        void init();
        /**
         * @brief 单帧推理函数，可覆写
         * 
         * @param frame 推理图像   
         * @param roi 
         * @return true 推理成功
         * @return false 推理失败
         */
        bool DeepDetectSingle(cv::Mat& frame, cv::Rect2d& roi);           // 模型推理检测主程序
        /**
         * @brief 双目推理（弃用）
         * 
         * @param left_frame 
         * @param right_frame 
         * @param left_roi 
         * @param right_roi 
         * @return true 
         * @return false 
         */
        bool DeepDetectDouble(cv::Mat &left_frame, cv::Mat &right_frame, cv::Rect2d &left_roi, cv::Rect2d &right_roi);
        /**
         * @brief opencv yaml读参（弃用
         * 
         * @param params 
         */
        // void setParam(cv::FileStorage fs);    
        /**
         * @brief Set the Param object
         * 
         * @param params 
         */
        void setParam(const DetectorParams params);                                         
        /**
         * @brief Debug输出函数
         * 
         */
        void DebugOutput();                                                        // Debug输出
        
        cv::Size2i getInputSize(); 
        
        private:                                                                                                                                                             
        bool yolov5_model_infer();                              // YOLOv5模型推理
        bool yolox_model_infer();                               // YOLOX模型推理
        // bool yolov8_model_infer();                               // YOLOv8模型推理

        /**
         * @brief 输入图像预处理
         * 
         * @param frame 
         * @param roi 
         * @return cv::Mat 
         */
        cv::Mat input_preprocess(cv::Mat& frame, cv::Rect2d& roi); // 输入图像处理
        /**
         * @brief 推理输出转Armor类
         * 
         * @param bboxes 推理输出结果bbox_t
         * @return Armors 
         */
        Armors bbox2Armor(std::vector<bbox_t> &bboxes);         


        public:     
        std::shared_ptr<InferenceBackendBase> m_backend; 
        DetectorParams m_params;
        //********** 外部参数 ***********//
        int Use_Deep;                       // 是否使用深度检测标志位
        std::string m_model_format = "trt"; // 模型所使用格式 OpenVINO转换IR格式(xml,bin) or ONNX格式
        // std::string m_xml;                  // 模型文件名称(不带文件类型后缀)
        std::string m_bin;
        std::string m_onnx;
        std::string m_model_type = "x";    // 模型类型，目前为 “v5” or “x”
        std::string DEVICE;         // 推理设备
        std::string m_model_file; 
        std::string m_input_precision; 
        bool m_debug;
        //********** 外部参数 ***********//
        /**
         * @brief 保存原图，可用于识别输出可视化（未使用
         * 
         */
        cv::Mat m_src;
        /**
         * @brief ROI信息
         * 
         */
        cv::Rect2d m_roi;
        cv::Mat input_mat;
        Armors m_armors;     // 检测出装甲板组
        //********** 后处理参数 ***********//
        int m_enemy_color;   // 敌方颜色
        int m_ncolor;      // 当前模型颜色数量
        int m_ntag;        // 当前模型类别数量
        int m_ntype;
        int input_height;              // 模型输入图像高度
        int input_width;               // 模型输入图像宽度
        float m_nmsThreshold;          // nms阈值
        float m_confThreshold;         // 置信度阈值
        //********** 后处理参数 ***********//

        int m_batch_size;

        float m_Lightratio = 4;        // 生成灯条比例常数

        private:
        int maxProposalCount;          // 模型输出通道数
        int objectLength;              // 模型输出单位长度
        int numberOfClass;             // 模型输出类别数
        float m_resize_scale;          // resize缩放尺寸
        std::map<int, int> m_label2id; // 模型输出类别序号与原类别映射表

        //********** YOLOX后处理参数 ***********//
        // std::vector<bbox_t> proposals;
        std::vector<bbox_t> m_bboxes;         // 模型预测bboxes
        std::vector<int> strides = {16, 32};
        std::vector<GridAndStride> grid_strides;
        //********** YOLOX后处理参数 ***********//

        void decode_YOLO_outputs(const float *prob, std::vector<bbox_t> &objects, const int img_w, const int img_h);
        // void decode_YOLOX_outputs(const float *prob, std::vector<bbox_t> &objects, const int img_w, const int img_h);
        void generate_grids_and_stride(const int target_w, const int target_h, std::vector<int>& strides, std::vector<GridAndStride>& grid_strides);       // 初始化grids_and_stride
        void generate_yolox_proposals(std::vector<GridAndStride> grid_strides, const float *feat_ptr, float prob_threshold, std::vector<bbox_t> &objects); // x版本生成预测
        void generate_yolov8_proposals(std::vector<GridAndStride> grid_strides, const float *feat_ptr, float prob_threshold, std::vector<bbox_t> &objects); // v8版本生成预测
        void nms_sorted_bboxes(const std::vector<bbox_t>& faceobjects, std::vector<int>& picked, float nms_threshold);                                     // 非极大值抑制

        int inputIndex,outputIndex;
        int input_size,output_size,gpu_img_size;
        float* m_prob;

    };
}


// Tools Functions


static float maxclass(const float *ptr, int nc, int &max_class)
{
    float max = 0;
    max_class = -1;
    for (int i = 0; i < nc; i++)
    {
        if (ptr[i] > max)
        {
            max_class = i;
            max = ptr[i];
        }
    }
    return max;
}

static void qsort_descent_inplace(std::vector<wmj::bbox_t> &faceobjects, int left, int right)
{
    int i = left;
    int j = right;
    float p = faceobjects[(left + right) / 2].conf;

    while (i <= j)
    {
        while (faceobjects[i].conf > p)
            i++;

        while (faceobjects[j].conf < p)
            j--;

        if (i <= j)
        {
            // swap
            std::swap(faceobjects[i], faceobjects[j]);

            i++;
            j--;
        }
    }

#pragma omp parallel sections
    {
#pragma omp section
        {
            if (left < j)
                qsort_descent_inplace(faceobjects, left, j);
        }
#pragma omp section
        {
            if (i < right)
                qsort_descent_inplace(faceobjects, i, right);
        }
    }
}

static void qsort_descent_inplace(std::vector<wmj::bbox_t> &objects)
{
    if (objects.empty())
        return;

    qsort_descent_inplace(objects, 0, objects.size() - 1);
}

static inline float intersection_area(const wmj::bbox_t &a, const wmj::bbox_t &b)
{
    cv::Rect_<float> inter = a.rect & b.rect;
    return inter.area();
}

template <typename _Tp> // 计算两点间的距离
double getDistance(const cv::Point_<_Tp> &p1, const cv::Point_<_Tp> &p2)
{

    const cv::Point_<_Tp> tmp = p1 - p2;

    return sqrt(tmp.ddot(tmp));
}